package scenario;

import org.junit.Test;

import java.util.NavigableSet;
import java.util.Objects;
import java.util.TreeSet;

// 思路1: 用数组表示座位
// 思路2: 用SortedMap存储已被占用的座位(BELOW)
// 思路3: 用堆保存两个已被占用座位间的距离
// 思路4: 用堆存储两个座位及距离(BELOW)
public class ExamRoom855 {

    public static class TestCase {
        @Test
        public void test1() {
            ExamRoom855 r = new ExamRoom855(10);
            printSeat(r);
            printSeat(r);
            printSeat(r);
            printSeat(r);
            printLeave(r, 4);
            printSeat(r);
        }

        @Test
        public void test2() {
            ExamRoom855 r = new ExamRoom855(8);
            printSeat(r);
            printSeat(r);
            printSeat(r);
            printLeave(r, 0);
            printLeave(r, 7);
            printSeat(r);
            printSeat(r);
            printSeat(r);
            printSeat(r);
            printSeat(r);
            printSeat(r);
            printSeat(r);
        }

        static void printSeat(ExamRoom855 r) {
            System.out.format("seat(): %d\n", r.seat());
        }

        static void printLeave(ExamRoom855 r, int p) {
            r.leave(p);
            System.out.format("leave(): %d\n", p);
        }
    }

    final int MIN = -1;
    int CAPACITY;
    SeatPair DUMMY_PAIR;
    NavigableSet<Integer> seatedSet = new TreeSet<>();
    NavigableSet<SeatPair> seatPairSet = new TreeSet<>();

    public ExamRoom855(int n) {
        this.CAPACITY = n;
        seatedSet.add(MIN);
        seatedSet.add(CAPACITY);
        DUMMY_PAIR = new SeatPair(MIN, CAPACITY);
        seatPairSet.add(DUMMY_PAIR);
    }

    public int seat() {
        if (seatedSet.size() >= CAPACITY +2) {
            return -1;
        }
        SeatPair p = seatPairSet.first();
        SeatPair p1 = new SeatPair(p.left, p.newSeat);
        SeatPair p2 = new SeatPair(p.newSeat, p.right);
        seatPairSet.remove(p);
        seatPairSet.add(p1);
        seatPairSet.add(p2);
        seatedSet.add(p.newSeat);
        return p.newSeat;
    }

    public void leave(int p) {
        if (p <= MIN || p >= CAPACITY || !seatedSet.contains(p)) {
            throw new RuntimeException("illegal seat p");
        }
        // 不能用floor或ceil, 它们包含等号
        int left = seatedSet.lower(p);
        int right = seatedSet.higher(p);
        SeatPair pair1 = new SeatPair(left, p);
        SeatPair pair2 = new SeatPair(p, right);
        seatPairSet.remove(pair1);
        seatPairSet.remove(pair2);
        seatedSet.remove(p);

        SeatPair pair3 = new SeatPair(left, right);
        seatPairSet.add(pair3);
    }

    class SeatPair implements Comparable<SeatPair> {
        public int left; // might be MIN
        public int right;// might be CAPACITY
        //public int distance;
        public int newSeat;
        public int newSeatDistance;

        public SeatPair(int left, int right) {
            this.left = left;
            this.right = right;
            chooseNewSeatAhead();
        }

        public void chooseNewSeatAhead() {
            if (left == MIN) {
                newSeat = 0;
                newSeatDistance = right - newSeat;
            } else if (right == CAPACITY) {
                newSeat = CAPACITY-1;
                newSeatDistance = newSeat - left;
            } else {
                newSeat = (left + right) / 2;
                newSeatDistance = newSeat - left;
            }
        }

        // 注意: distance=2 的可能要排在 distance=3 的前面, 因为left更小.
        // 由于两个虚拟端点的特殊性, 选择新seat的方案比较灵活. 所以不能用right-left来排序
        @Override
        public int compareTo(SeatPair that) {
            if (this.newSeatDistance != that.newSeatDistance) {
                return that.newSeatDistance - this.newSeatDistance;
            }
            return this.newSeat - that.newSeat;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            SeatPair seatPair = (SeatPair) o;
            return left == seatPair.left && right == seatPair.right;
        }

        @Override
        public int hashCode() {
            return Objects.hash(left, right);
        }
    }
}
